"""
The LossMSE class implements mean squared error loss.
"""

import numpy as np

class LossMSE:

    def __init__(self):
        """
        No initialization necessary.
        """
        pass

    def forward_pass(self, y_pred, y_true):
        """
        Compute the mean squared error loss.
        y_pred and y_true should be numpy arrays of shape (outputs, batch_size)
        """
        self.y_pred = y_pred
        self.y_true = y_true
        self.differences = y_pred - y_true
        self.squared_differences = self.differences ** 2
        return np.mean(self.squared_differences)

    def backward_pass(self):
        """
        Compute gradient of loss with respect to y_pred.
        """
        num_elements = np.prod(self.y_true.shape)
        self.dinputs = 2 * self.differences / num_elements
        return self.dinputs